import socket
import random
import pickle

# 生成私钥的函数
def generate_private_key(p):
    return random.randint(2, p - 2)

# 计算公钥
def generate_public_key(private_key, g, p):
    return pow(g, private_key, p)

# 计算共享密钥
def generate_shared_secret(private_key, other_public_key, p):
    return pow(other_public_key, private_key, p)

# 发送方的实现
def sender():
    # 公共素数 p 和基数 g
    p = 23  # 公共素数 p
    g = 5   # 基数 g

    # 生成发送方的私钥和公钥
    sender_private_key = generate_private_key(p)
    sender_public_key = generate_public_key(sender_private_key, g, p)

    print(f"B的私钥是: {sender_private_key}")
    print(f"B的公钥是: {sender_public_key}")

    # 设置接收方地址 (假设接收方在 localhost 上)
    receiver_host = 'localhost'
    receiver_port = 5000

    # 创建 socket 连接到接收方
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.connect((receiver_host, receiver_port))

        # 发送自己的公钥给接收方
        s.send(pickle.dumps(sender_public_key))

        # 接收接收方的公钥
        receiver_public_key = pickle.loads(s.recv(1024))

        print(f"A的公钥: {receiver_public_key}")

        # 计算共享密钥
        shared_secret = generate_shared_secret(sender_private_key, receiver_public_key, p)
        print(f"共享密钥是: {shared_secret}")
        key = shared_secret.to_bytes((shared_secret.bit_length() + 7) // 8, 'big')[:16]
        with open('shared_secret_B_path.txt', 'wb') as shared_secret_B_file:
            shared_secret_B_file.write(key)

# 启动发送方
if __name__ == "__main__":
    sender()
